package com.example.demo.algorithms;

import org.junit.Test;

import java.util.*;

/**
 * <p>Description: 找出n个数组中最大的k个元素</p>
 *
 * @author Eric Lee
 * @version v1.0.0
 * @since 2021/12/18 22:31
 **/
public class TopKInNArrays {

    int[][] data = {{5, 219, 405, 538, 845, 971}, {1000}, {4, 52, 99, 348, 691}};

    @Test
    public void test() {
        findTopK(data, 5);
    }

    /**
     * PriorityQueue 实现
     *
     * @param data
     * @param K
     */
    public void findTopK(int[][] data, int K) {
        PriorityQueue<Cell> maxHeap = new PriorityQueue<Cell>(new Comparator<Cell>() {
            public int compare(Cell num1, Cell num2) {
                return num2.val - num1.val;         // 最大堆比较器
            }
        });
        for (int i = 0; i < data.length; i++) {
            // 初始化堆
            int[] array = data[i];
            //记录是哪一个数组，索引位置和值
            Cell node = new Cell(i, array.length - 1, array[array.length - 1]);
            maxHeap.add(node);
        }
        // K大小不能超过数据总长度，注意索引是否越界

        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < K; i++) {
            Cell node = maxHeap.poll();
            sb.append(node.val + " ");
            int row = node.row;
            int index = node.index;
            if (index > 0) {
                Cell tmpNode = new Cell(row, index - 1, data[row][index - 1]);
                maxHeap.add(tmpNode);
            }
        }
        System.out.println(sb.toString().trim());
    }

    /**
     * 网络答案参考
     */
    @Test
    public void test2() {
        int k = 5;
        String res = getRes(data, k);
        System.out.println(res);
    }

    public static String getRes(int[][] arr, int k) {
        if (arr == null || arr.length == 0) return null;
        int nums = arr.length;
        Cell[] cells = new Cell[nums];
        for (int i = 0; i < nums; i++) {
            int val = arr[i][arr[i].length - 1];
            //记录是哪一个数组，索引位置和值
            Cell node = new Cell(i, arr[i].length - 1, val);
            cells[i] = node;
            heapInsert(cells, i);
        }
        int heapend = cells.length - 1;
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < k; i++) {
            Cell node = cells[0];
            sb.append(node.val + " ");
            int num = node.row;
            int index = node.index;
            if (index > 0) {
                cells[0] = new Cell(num, index - 1, arr[num][index - 1]);
                heapfy(cells, heapend);
            } else {
                if (heapend <= 0) {
                    throw new RuntimeException("数不够");
                }
                cells[0] = cells[heapend];
                heapfy(cells, --heapend);
            }
        }
        return sb.toString().trim();
    }

    public static void heapfy(Cell[] cells, int heapend) {
        int left = 1;
        int right = 2;
        int index = 0;
        int largest = index;
        while (left <= heapend) {
            if (cells[left].val > cells[index].val) {
                largest = left;
            }
            if (right <= heapend && cells[right].val > cells[largest].val) {
                largest = right;
            }
            if (largest != index) {
                swap(cells, index, largest);
            } else {
                break;
            }
            left = largest * 2 + 1;
            right = largest * 2 + 2;
        }
    }

    public static void heapInsert(Cell[] cells, int index) {
        while (index > 0) {
            int parent = (index - 1) / 2;
            if (cells[parent].val < cells[index].val) {
                swap(cells, parent, index);
                index = parent;
            } else {
                break;
            }
        }
    }

    public static void swap(Cell[] cells, int i, int j) {
        Cell tmp = cells[i];
        cells[i] = cells[j];
        cells[j] = tmp;
    }

    static class Cell {
        int row;
        int index;
        int val;

        public Cell(int row, int index, int val) {
            this.row = row;
            this.index = index;
            this.val = val;
        }
    }
}
